// this set of kernels compiles successfully at this point
// we need to figure out which data structures to pass in
// from the algorithm
//
//
//#include "/sa/development/mpi_lasso/src/ssvs_dimensions.h"
#include "/home/garykche/gary/code/src/ssvs_dimensions.h"
//#include "/export/home/garykche/development/imputation/ayers/cl_constants.h"

__kernel void simple(
__constant int * scaler,
__global int * outputvec
) {
  int threadindex = get_local_id(0);
  outputvec[threadindex] = *scaler * threadindex;
  return;
}

__kernel void xprimex(
__const int models,
__const int samplesize,
__const int variables,
__const int subsetsize,
__const int subsetsize2,
__global int * subset,
__global float * designmat,
__global float * xx,
__global float * inv_xx,
__local float * local_sourceblock,
__local int * local_subset,
//__local float * local_vec1,
__local float * local_vec2,
__local float * local_xx,
__local float * local_chol,
__local float * local_inv_xx
) {
  int model = get_group_id(0);
  int threadindex = get_local_id(0);
  //local_vec1[threadindex] = 0;
  local_vec2[threadindex] = 0;
  if (threadindex<subsetsize) 
  local_subset[threadindex] = subset[model*subsetsize+threadindex];
  for(int chunk=0;chunk<subsetsize2/BLOCK_WIDTH+1;++chunk){
    int element = chunk*BLOCK_WIDTH+threadindex;
    if (element<subsetsize2) local_xx[element] = 0;
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  for(int chunk=0;chunk<(samplesize/BLOCK_WIDTH)+1;++chunk){
    int person = chunk*BLOCK_WIDTH + threadindex;
    for(int j=0;j<subsetsize;++j){
      local_sourceblock[j*BLOCK_WIDTH+threadindex] = person<samplesize?
      designmat[local_subset[j]*samplesize+person]:0;
    }
    //if (person>=samplesize){
      //local_vec1[threadindex] = 0;
      //local_vec2[threadindex] = 0;
    //}
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int j=0;j<subsetsize;++j){
      //if (person<samplesize) local_vec1[threadindex] = 
      //designmat[local_subset[j]*samplesize+person];
      //designmat[subset[model*subsetsize+j]*samplesize+person];
      //1;
      //barrier(CLK_LOCAL_MEM_FENCE);
      for(int k=0;k<subsetsize;++k){
        //if (person<samplesize) 
        local_vec2[threadindex] = local_sourceblock[j*BLOCK_WIDTH+threadindex]*local_sourceblock[k*BLOCK_WIDTH+threadindex];
        //local_vec2[threadindex] = local_vec1[threadindex] * 
        //1;
        //designmat[local_subset[k]*samplesize+person];
        //designmat[subset[model*subsetsize+k]*samplesize+person];
        barrier(CLK_LOCAL_MEM_FENCE);
        for(int s=BLOCK_WIDTH/2; s>0; s>>=1) {
          if (threadindex<s) {
             local_vec2[threadindex]+=local_vec2[threadindex+s];
          }
          barrier(CLK_LOCAL_MEM_FENCE);
        }
        if (threadindex==0){
          //local_xx[j*subsetsize+k] += 1;
          local_xx[j*subsetsize+k] += local_vec2[0];
        }
      }
    }
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  for(int chunk=0;chunk<subsetsize2/BLOCK_WIDTH+1;++chunk){
    int element = chunk*BLOCK_WIDTH+threadindex;
    if (element<subsetsize2) {
       local_chol[element] = local_xx[element];
       //xx[model*subsetsize2+element] = local_xx[element];
    }
  }
  for(int i=0;i<subsetsize;++i){
    if (threadindex<subsetsize){
      local_inv_xx[i*subsetsize+threadindex] = (threadindex==i);
    }
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  // DONE INITIALIZING DATA STRUCTURES FOR CHOLESKY DECOMP INVERSION
  if (threadindex==0){
    for(int i=0;i<subsetsize;++i){
      for(int j=i;j<subsetsize;++j){
        float sum = local_xx[i*subsetsize+j];
        for(int k=0;k<i;++k){
          sum-=local_chol[i*subsetsize+k]*local_chol[j*subsetsize+k];
        }
        local_chol[j*subsetsize+i] = (i==j)?sqrt(sum):sum/local_chol[i*subsetsize+i];
      }
    }
    for(int i=0;i<subsetsize;++i){
      for(int j=0;j<subsetsize;++j){
        float sum = local_inv_xx[i*subsetsize+j];
        for(int k=0;k<j;++k){
          sum-=local_chol[j*subsetsize+k]*local_inv_xx[i*subsetsize+k];
        }
        local_inv_xx[i*subsetsize+j] = sum/local_chol[j*subsetsize+j];
      }
      for(int j=subsetsize-1;j>=0;--j){
        float sum = local_inv_xx[i*subsetsize+j];
        for(int k=subsetsize-1;k>j;--k){
          sum-=local_chol[k*subsetsize+j]*local_inv_xx[i*subsetsize+k];
        }
        local_inv_xx[i*subsetsize+j] = sum/local_chol[j*subsetsize+j];
      }
    }
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  for(int chunk=0;chunk<subsetsize2/BLOCK_WIDTH+1;++chunk){
    int element = chunk*BLOCK_WIDTH+threadindex;
    if (element<subsetsize2) inv_xx[model*subsetsize2+element] = local_inv_xx[element];
  }
  return;
}

__kernel void invert(
__const int models,
__const int subsetsize,
__const int subsetsize2,
__global float * xx,
__global float * inv_xx,
__local float * local_xx,
__local float * local_chol,
__local float * local_inv_xx
) {
  int model = get_group_id(0);
  int threadindex = get_local_id(0);
  for(int chunk=0;chunk<((subsetsize2)/SMALL_BLOCK_WIDTH)+1;++chunk){
    int element = chunk*SMALL_BLOCK_WIDTH+threadindex;
    if (element<subsetsize2){
       local_xx[element] = xx[model*subsetsize2+element];
       local_chol[element] = local_xx[element];
    }
  }
  for(int i=0;i<subsetsize;++i){
    if (threadindex<subsetsize){
      local_inv_xx[i*subsetsize+threadindex] = (threadindex==i);
    }
  }
  if (threadindex==0){
    for(int i=0;i<subsetsize;++i){
      for(int j=i;j<subsetsize;++j){
        float sum = local_xx[i*subsetsize+j];
        for(int k=0;k<i;++k){
          sum-=local_chol[i*subsetsize+k]*local_chol[j*subsetsize+k];
        }
        local_chol[j*subsetsize+i] = (i==j)?sqrt(sum):sum/local_chol[i*subsetsize+i];
      }
    }
    for(int i=0;i<subsetsize;++i){
      for(int j=0;j<subsetsize;++j){
        float sum = local_inv_xx[i*subsetsize+j];
        for(int k=0;k<j;++k){
          sum-=local_chol[j*subsetsize+k]*local_inv_xx[i*subsetsize+k];
        }
        local_inv_xx[i*subsetsize+j] = sum/local_chol[j*subsetsize+j];
      }
      for(int j=subsetsize-1;j>=0;--j){
        float sum = local_inv_xx[i*subsetsize+j];
        for(int k=subsetsize-1;k>j;--k){
          sum-=local_chol[k*subsetsize+j]*local_inv_xx[i*subsetsize+k];
        }
        local_inv_xx[i*subsetsize+j] = sum/local_chol[j*subsetsize+j];
      }
    }
  }
  for(int chunk=0;chunk<subsetsize2/SMALL_BLOCK_WIDTH+1;++chunk){
    int element = chunk*SMALL_BLOCK_WIDTH+threadindex;
    if (element<subsetsize2) inv_xx[model*subsetsize2+element] = local_inv_xx[element];
  }
  return;
}

__kernel void compute_xxx(
__const int models,
__const int samplesize,
__const int subsetsize,
__const int subsetsize2,
__global int * subset,
__global float * designmat,
__global float * inv_xx,
__global float * xxx,
__local float * local_inv_xx,
__local float * local_sourceblock,
__local float * local_destblock
) {
  int model = get_group_id(0);
  int threadindex = get_local_id(0);
  //for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
  //  for(int j=0;j<subsetsize;++j){
  //    int personindex = personchunk*BLOCK_WIDTH+threadindex;
  //    if (personindex<samplesize){
  //      xxx[model*(subsetsize*samplesize)+j*samplesize+personindex] =  0;
  //    }
  //  }
  //}
  //return;
  for(int chunk=0;chunk<subsetsize2/BLOCK_WIDTH+1;++chunk){
    int element = chunk*BLOCK_WIDTH+threadindex;
    if (element<subsetsize2) local_inv_xx[element] = inv_xx[model*subsetsize2+element];
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
    for(int j=0;j<subsetsize;++j){
      int personindex = personchunk*BLOCK_WIDTH+threadindex;
      local_destblock[j*BLOCK_WIDTH+threadindex] = 0;
      local_sourceblock[j*BLOCK_WIDTH+threadindex] = personindex<samplesize?designmat[subset[model*subsetsize+j]*samplesize+personindex]:0;

    }
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int row=0;row<subsetsize;++row){
      for(int col=0;col<subsetsize;++col){
        local_destblock[row*BLOCK_WIDTH+threadindex]+=local_inv_xx[row*subsetsize+col] * local_sourceblock[col*BLOCK_WIDTH+threadindex];

      }
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int j=0;j<subsetsize;++j){
      int personindex = personchunk*BLOCK_WIDTH+threadindex;
      if (personindex<samplesize){
        //xxx[model*(subsetsize*samplesize)+j*samplesize+personindex] =  model;
        xxx[model*(subsetsize*samplesize)+j*samplesize+personindex] =  local_destblock[j*BLOCK_WIDTH+threadindex];
      }
    }
  }
  return;
}

__kernel void compute_beta(
__const int models,
__const int samplesize,
__const int subsetsize,
__const int subsetsize2,
__global int * subset,
__global float * designmat,
__global float * inv_xx,
__global float * trait,
__global float * beta,
__local int * local_subset,
__local float * local_trait, // SAMPLESIZE
__local float * local_inv_xx,
__local float * local_sourceblock,
__local float * local_destblock,
__local float * local_beta_reduction,
__local float * local_beta
) {
  int model = get_group_id(0);
  int threadindex = get_local_id(0);
  for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
    int personindex = personchunk*BLOCK_WIDTH+threadindex;
    if(personindex<samplesize) local_trait[personindex] = 
    trait[personindex];
  }
  if (threadindex<subsetsize) {
     local_beta[threadindex] = 0;
     local_subset[threadindex] = subset[model*subsetsize+threadindex];
  }
  for(int chunk=0;chunk<subsetsize2/BLOCK_WIDTH+1;++chunk){
    int element = chunk*BLOCK_WIDTH+threadindex;
    if (element<subsetsize2) local_inv_xx[element] = inv_xx[model*subsetsize2+element];
  }
  for(int j=0;j<subsetsize;++j){
    local_beta_reduction[j*BLOCK_WIDTH+threadindex] = 0;
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  // INITIALIZED LOCAL ARRAYS
  for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
    int personindex = personchunk*BLOCK_WIDTH+threadindex;
    for(int j=0;j<subsetsize;++j){
      local_destblock[j*BLOCK_WIDTH+threadindex] = 0;
      //local_sourceblock[j*BLOCK_WIDTH+threadindex] = personindex<samplesize?designmat[subset[model*subsetsize+j]*samplesize+personindex]:0;
      local_sourceblock[j*BLOCK_WIDTH+threadindex] = personindex<samplesize?designmat[local_subset[j]*samplesize+personindex]:0;

    }
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int row=0;row<subsetsize;++row){
      for(int col=0;col<subsetsize;++col){
        local_destblock[row*BLOCK_WIDTH+threadindex]+=local_inv_xx[row*subsetsize+col] * local_sourceblock[col*BLOCK_WIDTH+threadindex];

      }
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    for(int j=0;j<subsetsize;++j){
      local_beta_reduction[j*BLOCK_WIDTH+threadindex] += personindex<samplesize?local_destblock[j*BLOCK_WIDTH+threadindex]*local_trait[personindex]:0;
    }
    barrier(CLK_LOCAL_MEM_FENCE);
  }
  for(int j=0;j<subsetsize;++j){
    for(int s=BLOCK_WIDTH/2; s>0; s>>=1) {
      if (threadindex<s) {
        local_beta_reduction[j*BLOCK_WIDTH+threadindex]+=local_beta_reduction[j*BLOCK_WIDTH+threadindex+s];
      }
      barrier(CLK_LOCAL_MEM_FENCE);
    }
    if(threadindex==0) local_beta[j] = local_beta_reduction[j*BLOCK_WIDTH];
  }
  if (threadindex<subsetsize) beta[model*subsetsize+threadindex] = local_beta[threadindex];
  return;
}

__kernel void compute_beta_old(
__const int models,
__const int samplesize,
__const int subsetsize,
__const int subsetsize2,
__global float * xxx,
__global float * trait,
__global float * beta,
__local float * local_trait, // SAMPLESIZE
__local float * local_xxx, // BLOCK_WIDTH
__local float * local_beta
) {
  int model = get_group_id(0);
  int threadindex = get_local_id(0);
  for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
    int personindex = personchunk*BLOCK_WIDTH+threadindex;
    if(personindex<samplesize) local_trait[personindex] = 
    trait[personindex];
  }
  if (threadindex<subsetsize) local_beta[threadindex] = 0;
  local_xxx[threadindex] = 0;
  barrier(CLK_LOCAL_MEM_FENCE);
  for(int j=0;j<subsetsize;++j){
    for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
      int personindex = personchunk*BLOCK_WIDTH+threadindex;
      local_xxx[threadindex] = (personindex<samplesize)?
      xxx[model*(subsetsize*samplesize)+j*samplesize+personindex] * local_trait[personindex]:0; 
      barrier(CLK_LOCAL_MEM_FENCE);
      for(int s=BLOCK_WIDTH/2; s>0; s>>=1) {
        if (threadindex<s) {
           local_xxx[threadindex]+=local_xxx[threadindex+s];
        }
        barrier(CLK_LOCAL_MEM_FENCE);
      }
      if(threadindex==0) local_beta[j] += local_xxx[threadindex];
    }
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  if (threadindex<subsetsize) beta[model*subsetsize+threadindex] = local_beta[threadindex];
  return;
}


__kernel void compute_rsq(
__const int models,
__const int samplesize,
__const int subsetsize,
__const int subsetsize2,
__const float trait_mean,
__const float tss,
__global int * subset,
__global float * beta,
__global float * designmat,
__global float * rsq,
__global float * trait,
__local int * local_subset,
__local float * local_trait, // SAMPLESIZE
__local float * local_fitted, // SAMPLESIZE
__local float * local_sse, // SMALL_BLOCK_WIDTH
__local float * local_beta
) {
  int model = get_group_id(0);
  int threadindex = get_local_id(0);
  if (threadindex<subsetsize){
     local_beta[threadindex] = beta[model*subsetsize+threadindex];
     local_subset[threadindex] = subset[model*subsetsize+threadindex];
  }
  for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
    int personindex = personchunk*BLOCK_WIDTH+threadindex;
    if(personindex<samplesize) {
      local_fitted[personindex] = 0;
      local_trait[personindex] = trait[personindex];
    }
  }
  local_sse[threadindex] = 0;
  barrier(CLK_LOCAL_MEM_FENCE);
  for(int j=0;j<subsetsize;++j){
    for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
      int personindex = personchunk*BLOCK_WIDTH+threadindex;
      if (personindex<samplesize){
        local_fitted[personindex] += local_beta[j]*designmat[local_subset[j]*samplesize+personindex];
        //local_fitted[personindex] += local_beta[j]*designmat[subset[model*subsetsize+j]*samplesize+personindex];
      }
      barrier(CLK_LOCAL_MEM_FENCE);
    }
  }
  for(int personchunk=0;personchunk<samplesize/BLOCK_WIDTH+1;++personchunk){
    int personindex = personchunk*BLOCK_WIDTH+threadindex;
    if(personindex<samplesize) {
      local_sse[threadindex]+=pow(local_trait[personindex]-local_fitted[personindex],2);
      //local_sse[threadindex]+=pow(local_fitted[personindex]-trait_mean,2);
    }
    barrier(CLK_LOCAL_MEM_FENCE);
  }
  for(int s=BLOCK_WIDTH/2; s>0; s>>=1) {
    if (threadindex<s) {
      local_sse[threadindex]+=local_sse[threadindex+s];
    }
    barrier(CLK_LOCAL_MEM_FENCE);
  }
  if(threadindex==0){
    //rsq[threadindex] = local_sse[threadindex];
    rsq[model] = -.5*samplesize*(1.837877+log(1.*local_sse[threadindex]/(samplesize-subsetsize))+1);
threadindex;
    //rsq[model] = 1-local_sse[threadindex]/tss;
  }
  return;
}
